(我故意把原版的放著 然後就拖到現在的進度了 等我重賽版的上完 再放回去原本的版本ㄅ)

昨天我們把 A8 @ B8 拆成最笨的 naive INT8 GEMM。

一個 thread 算一個 C[i,j]，for loop 從 k = 0 跑到最後。這種寫法很好懂，也很好對答案，但它沒有用到今天真正想用的硬體。

今天要換問題：

怎樣才算真的走到 INT8 Tensor Core？

答案不是「資料型別用了 int8_t 就會自動變快」。

如果你只是寫：

sum += (int32_t)a8 * (int32_t)b8;

compiler 不會自動把它變成 Tensor Core GEMM。它比較像一般 CUDA core 上的整數乘加。要真的吃到 Tensor Core，我們要把資料排成硬體想要的 tile，讓一個 warp 發出 mma.sync。

今天會第一次碰到 cp.async、ldmatrix、mma.sync。但先說清楚：這篇不是 PTX 教學，也不是完整 assembly walkthrough。今天只要把 Phase 22 的 fast kernel 骨架講懂：

global memory 裡的 INT8 residue tile
    ↓ cp.async
shared memory
    ↓ ldmatrix
register fragment
    ↓ mma.sync
int32 accumulator
    ↓ modulo + CRT
double C

這仍然是 Phase 22。也就是 double input/output、約 34-bit mantissa 的 fast path。走上 Tensor Core 只代表 INT8 GEMM 變快，不代表精度突然變成 full FP64。

先記得昨天的資料長什麼樣

Phase 22 的資料不是原始 double。

它先把每個 double 拆成 high / low：

hi = round(v × 2^17)
lo = round((v - hi / 2^17) × 2^34)

然後每個 slice 再對 7 個質數取餘數：

127, 113, 109, 107, 103, 101, 97

所以 A 會有：

A_hi_mod[p]
A_lo_mod[p]

B 也會有：

B_hi_mod[p]
B_lo_mod[p]

真正要算的不是一次 GEMM，而是四個 Ozaki pass：

A_hi × B_hi
A_hi × B_lo
A_lo × B_hi
A_lo × B_lo

每個 pass 裡又要跑 7 個質數。

所以很粗地看，Phase 22 fast path 有：

4 個 high/low pass × 7 個 prime = 28 個 INT8 GEMM 的工作量

這就是為什麼 Day 8 的 naive GEMM 只能拿來對答案。28 次 naive GEMM 一定慢到不能當主線。

Tensor Core 要的是 tile，不是普通 for loop

Day 3 的 GEMM 可能長這樣：

C[row, col] = Σ A[row, kk] × B[kk, col]

這是數學視角。

Tensor Core 的視角不是這樣。它要的是一小塊矩陣乘法，一次由一個 warp 合作完成。例如我們這條路最後用的是類似：

mma.sync.m16n8k32.s8

意思大概是：

A fragment: 16 × 32, int8
B fragment: 32 × 8,  int8
C fragment: 16 × 8,  int32 accumulator

一次 mma.sync 不是一個 thread 做一個乘法，而是一整個 warp 合作，把一小塊矩陣乘積加到 accumulator 裡。

所以 host 端看到的 kernel launch 還是一顆 CUDA kernel，但 kernel 裡的基本單位已經變成：

thread block
    → warp
        → mma fragment

不是：

thread
    → scalar multiply-add

Phase 22 的 tile 形狀

Phase 22 fast path 目前用的 C tile 是：

128 × 64

K 方向每次吃：

32

可以把一個 CTA 想成負責這塊：

C tile: 128 rows × 64 cols

每次從 global memory 搬：

A tile: 128 × 32
B tile:  64 × 32

然後用 8 個 warp 去算這個 C tile。每個 warp 再負責其中一塊小 tile。實作上會反覆發出 m16n8k32 的 MMA，把小片小片的結果加到 int32 accumulator 裡。

用圖看比較像：

          B tile, one prime
          64 cols × 32 K
        ┌────────────────┐
        │                │
        │                │
        └────────────────┘

A tile  ┌────────────────┐
128 rows│                │   →   C tile 128 × 64
× 32 K  │                │
        │                │
        │                │
        └────────────────┘

這裡的 A/B tile 都是某個 prime 底下的 residue，不是原本的 double。

cp.async：先把 global memory 搬到 shared memory

Tensor Core 算很快。真的跑起來時，常常不是乘法本身最麻煩，而是資料怎麼餵進去。

如果每次 MMA 都直接從 global memory 抓 A/B，速度會被記憶體延遲拖死。所以 fast kernel 會先把 A/B tile 搬進 shared memory。

在 Ampere/Ada 這類卡上，可以用 cp.async 做非同步搬運：

global memory → shared memory

它的好處是：當現在這一塊 tile 正在算的時候，可以先把下一塊 K tile 搬進 shared memory。這就是雙緩衝的想法：

buffer 0 正在算
buffer 1 正在載入下一塊

下一輪交換：

buffer 1 正在算
buffer 0 正在載入下一塊

所以 kernel 裡常見的節奏會像：

load first K tile
wait

for each K tile:
    prefetch next K tile
    ldmatrix current tile
    mma current tile
    wait next tile

這裡的 wait 和 __syncthreads() 不能亂放。少了會讀到還沒搬完的 shared memory；放太多又會把 pipeline 打斷。

Day 9 的重點不是把同步優化到最好，而是先確認：

資料有正確搬進 shared memory
MMA 讀到的是正確 tile
accumulator 對得上 Day 8 naive path

ldmatrix：shared memory 到 register fragment

資料進 shared memory 後，還不能直接丟給 mma.sync。

mma.sync 要吃的是 register fragment。也就是每個 thread 持有某些 A/B fragment 的一部分，整個 warp 合起來剛好組出一個 MMA 需要的小矩陣。

這一步會用 ldmatrix：

shared memory → register fragment

這也是為什麼 Day 8 說 A/B 的 layout 很重要。

naive GEMM 可以很單純：

A8[row, kk]
B8[kk, col]

但 ldmatrix 對 shared memory 的排列方式有自己的期待。如果 layout 排錯，程式不一定 crash，但 MMA 會拿錯元素相乘。最後看起來就像數學錯了，其實只是 memory layout 錯。

所以開發時我會把這幾層分開驗：

global memory 的 A8/B8 對不對
shared memory 裡的 tile 對不對
ldmatrix 讀出的 fragment 對不對
mma accumulator 對不對

不要一開始就看最後 double 的 max error。

mma.sync：真正的 INT8 GEMM

到了 mma.sync 這一步，才是真的 Tensor Core。

概念上就是：

acc += A_fragment × B_fragment

其中：

A_fragment: int8
B_fragment: int8
acc:        int32

對單一 prime 來說，kernel 會沿著 K 方向一直累加：

for kk in K tiles:
    load A_p tile
    load B_p tile
    acc_p += A_p × B_p

跑完整個 K 以後，acc_p 代表：

Σ A_p[i,k] × B_p[k,j]

也就是這個 prime 底下的 GEMM accumulator。

接著才做：

acc_p mod p

因為 CRT 要的是每個質數世界裡的答案：

(A × B) mod p

不是沒有取 mod 的 int32 sum。

7 個質數目前還是序列

一個很自然的問題是：

既然有 7 個 prime，能不能一次一起算？

Phase 22 fast path 的答案是：沒有真正一次全部融合成一個超大的 MMA 問題。它比較像：

for p in primes:
    對這個 p 跑完整個 K 的 INT8 MMA
    acc_p = acc_p mod p
    weighted_sum += acc_p × f_p

也就是 7 個 prime 仍然是序列處理。

這不是數學上不能平行，而是工程上的取捨。每個 prime 都有自己的 A/B residue 矩陣，也有自己的 accumulator。你如果想把 7 個 prime 的 accumulator 全部同時留在 register 裡，register pressure 會非常高。留不住就會 spill 到 local memory，速度會掉得很慘。

另一個做法是把每個 prime 的 C_p 都寫回 global memory：

C_127
C_113
...
C_97

然後再開一顆 kernel 做 CRT。

這樣正確性很好查，但 HBM 讀寫會變多。Phase 22 fast path 選的是另一個方向：在一顆 kernel 裡，MMA 做完就直接把 acc_p 拿來 CRT，不把 7 張 C_p 全部落地。

這就是 fused 的第一個意義：

不是把所有事情變成一行魔法
而是避免中間矩陣一直寫回 global memory

CRT 在 accumulator 後面接上來

每個 prime 的 MMA 完成後，會得到 int32 accumulator。

接著做：

r_p = acc_p mod p

再乘上 Day 7 的重建係數：

weighted_sum += r_p × f_p
weighted_sum %= M

7 個 prime 都處理完後：

x = weighted_sum
if x > M/2:
    x -= M

這個 x 就是某個 high/low pass 的整數內積。

最後才乘上該 pass 的 scale：

C[i,j] += x × inv

四個 pass 的 inv 分別是：

hi × hi  →  2^-34
hi × low →  2^-51
low × hi →  2^-51
low × low→  2^-68

所以整個 Phase 22 的 fast path 可以想成：

for pass in {HH, HL, LH, LL}:
    for p in 7 primes:
        int8 mma over full K
        r_p = acc_p mod p
        CRT accumulate
    signed reconstruct
    C += reconstruct × pass_scale

這樣講比較慢，但它比直接說「fused CRT kernel」清楚很多。

為什麼 17-bit slice 需要 CRT

這裡要補一個容易混淆的點。

INT8 Tensor Core 一次只能吃 s8 × s8 → s32。那為什麼我們不直接把 hi 當 int8 slice 算？

因為 Phase 22 的 high slice 是：

hi = round(v × 2^17)

如果 v 在 [-1, 1]，hi 可以接近：

±131072

這不是 int8。

所以我們才把 hi 再轉成：

hi mod 127
hi mod 113
...

每個 residue 才塞得進目前 Phase 22 的 int8_t 資料路徑。

這跟 reference 裡的 Scheme I 不一樣。Scheme I 通常會把 mantissa 切成很多個真正的 8-bit slice，讓每個 slice pair 的乘積和 K 累加能放進 int32，然後用正確位權加回去。那條路可以少掉 Phase 22 這種「每個 17-bit slice pair 還要再跑 7 個 prime」的 CRT 成本。

所以兩條路的成本不一樣：

Phase 22:
    2 個 17-bit slices
    4 個 high/low pass
    每個 pass 7 個 prime
    → 28 份 INT8 GEMM 工作

Scheme I reference:
    多個 8-bit slices
    多個 slice pair
    通常不需要每個 pair 再套 7-prime CRT

這也是為什麼我們現在可以贏 native FP64，但同一個 34-bit 精度預算下，還是輸 cuBLAS emulation。不是因為 CRT 數學錯，而是這條資料路徑比較重。

這不是 full FP64

再重複一次，因為這件事很容易被 double 這個型別騙到。

Phase 22 的 kernel 最後寫回的是 double C。但它重建的是：

17-bit high + 17-bit low

也就是約 34-bit mantissa 的近似。走上 mma.sync 只會讓這個 34-bit-ish 計算變快，不會讓它變成 53-bit FP64。

如果要往 full FP64 靠近，需要另一組設計：

更多 mantissa slices
更多 pair schedule
或更多 CRT moduli
完整 scaling / undo scaling
必要時 auto mode / fallback

這些會在後面講 reference 和修正版規劃時再展開。Day 9 的任務很窄：

把 Day 8 對過答案的 INT8 GEMM
搬到 Tensor Core 上

怎麼對答案

Day 9 最怕的是：kernel 跑很快，但錯在哪裡完全看不出來。

所以驗證順序還是延續 Day 8：

第一層：對單一 prime 的 accumulator

先固定：

p = 127
pass = high × high

讓 Tensor Core kernel 算出：

acc_127[i,j]

然後跟 Day 8 naive INT8 GEMM 的 accumulator 對。

這一步不要做 CRT，也不要乘 scale。

如果這裡錯，通常是：

shared memory layout
ldmatrix mapping
warp tile index
K tail / padding

第二層：對 acc mod p

accumulator 對了，再看：

r_127 = acc_127 mod 127

如果 accumulator 對，但 remainder 錯，才去查 % p 或 magic division。

第三層：對 7-prime CRT

把 7 個 prime 的 residue 都湊齊後，才做：

x = CRT(r_127, r_113, ..., r_97)

這一步要刻意測負數，確認 signed fold 有過：

if x > M/2:
    x -= M

第四層：對四個 pass 的 double

最後才把四個 pass 加回去：

C = HH × 2^-34
  + HL × 2^-51
  + LH × 2^-51
  + LL × 2^-68

這時才看 max error。

目前 Phase 22 在 2048³、U(-1,1) 上會看到大約：

max error ≈ 1e-9

這是符合 34-bit-ish 設計的結果。不要期待 1e-15。

這顆 kernel 為什麼還不夠快

走上 Tensor Core 之後，不代表問題結束。

Day 9 的骨架能跑，但還有很多東西會吃掉效能：

1. 每個 prime 都要掃完整個 K
2. modulo p 和 modulo M 都不是免費
3. shared memory layout 可能有 bank conflict
4. register pressure 會限制 occupancy
5. 7 個 prime 的中間結果不能全部舒服地留在 register

這就是 Day 10 要處理的問題。

Day 9 的成功標準不是「吃滿 Tensor Core」。今天只要做到：

Tensor Core path 能跑
結果跟 naive path 對得上
end-to-end error 維持在 ~1e-9

真的要追速度，要到下一篇才開始看 % p、Barrett、swizzle、bank conflict、occupancy。

工程筆記：不要被 mma.sync 騙了

mma.sync 出現在程式裡，不代表 kernel 就一定快。

我一開始也很容易掉進這個陷阱：看到 Tensor Core 指令出現，就以為瓶頸解了。實際上不是。

Tensor Core 只是算力單元。你還要把資料用正確節奏餵給它：

global memory load 要 coalesced
shared memory layout 要配合 ldmatrix
cp.async pipeline 不能斷
register 不能爆
occupancy 不能太低
CRT 的整數運算不能蓋過 MMA

所以今天先把 kernel 分成兩個評價：

正確性：有沒有跟 naive path 對上？
效能：Tensor pipe 有沒有真的吃起來？

Day 9 只把第一件事收斂。

Day 10 才開始處理第二件事。

今天的產出

Day 9 結束時，應該有一顆 crt_fast_kernel 的骨架：

1. CTA 算 128×64 的 C tile。
2. K 方向以 32 為單位搬 A/B residue tile。
3. cp.async 把 global memory 搬到 shared memory。
4. ldmatrix 把 shared memory 讀成 register fragment。
5. mma.sync.m16n8k32.s8 做 INT8 GEMM。
6. 每個 prime 的 accumulator 做 mod p。
7. 用 CRT 係數拼回整數。
8. 乘上 high/low pass 的 scale，加回 double C。

這時候它已經不是 naive kernel 了。

但也還不是最終優化版。

今天先停在這裡

今天可以帶走三句：

1. int8 不等於 Tensor Core。 要真的走 Tensor Core，需要 tile、shared memory、ldmatrix、mma.sync 全部配合。
2. Phase 22 的 7 個 prime 仍然是序列成本。 每個 prime 都要做一輪 INT8 MMA，做完才接 CRT。
3. mma.sync 只解決速度，不解決精度。 Phase 22 仍然是約 34-bit mantissa 的 fast path，不是 full FP64。

明天開始看為什麼這顆 kernel 還沒有吃滿：% p、% M、shared memory bank conflict、occupancy，以及為什麼跟 cuBLAS emulation 還差一段距離。